"""
Foldseek结构搜索工具
调用后端API /api/search/structure/ 进行Foldseek结构搜索
"""

import json
import requests
from qwen_agent.tools import BaseTool
from typing import Dict, Any
import os
import config


class FoldseekStructureSearchTool(BaseTool):
    """Foldseek结构搜索工具类"""
    
    # 必须设置的类属性
    name = "foldseek_structure_search"
    description = "使用Foldseek算法搜索与给定蛋白质3D结构相似的蛋白质。支持从数据库选择已有蛋白质的PDB文件或上传自定义PDB文件。"
    parameters = {
        "type": "object",
        "properties": {
            "query_accession": {
                "type": "string",
                "description": "数据库中蛋白质的accession（例如'WP_042059674.1'），用于从数据库选择PDB文件进行分析"
            },
            "pdb_file_path": {
                "type": "string",
                "description": "本地PDB文件的绝对路径，用于上传自定义PDB文件进行分析"
            },
            "tm_score_threshold": {
                "type": "number",
                "description": "TM-score阈值，用于过滤结果（默认0.5，范围0-1）",
                "default": 0.5
            },
            "max_results": {
                "type": "integer",
                "description": "最大返回结果数（默认100，范围10-500）",
                "default": 100
            }
        },
        "required": []  # 需要提供query_accession或pdb_file_path其中之一
    }
    
    def call(self, params: Dict[str, Any], **kwargs) -> str:
        """
        执行Foldseek结构搜索
        
        Args:
            params: 参数字典，包含query_accession或pdb_file_path、tm_score_threshold、max_results
            **kwargs: 其他关键字参数
            
        Returns:
            str: JSON格式的字符串结果
        """
        try:
            # 参数解析和验证
            if isinstance(params, str):
                params = json.loads(params)
            
            query_accession = params.get("query_accession", "").strip()
            pdb_file_path = params.get("pdb_file_path", "").strip()
            
            # 验证输入：必须提供其中之一
            if not query_accession and not pdb_file_path:
                return json.dumps({
                    "success": False,
                    "error": "必须提供query_accession或pdb_file_path其中之一"
                }, ensure_ascii=False)
            
            if query_accession and pdb_file_path:
                return json.dumps({
                    "success": False,
                    "error": "只能选择一种查询方式（query_accession或pdb_file_path）"
                }, ensure_ascii=False)
            
            # 如果使用pdb_file_path，验证文件是否存在
            if pdb_file_path:
                if not os.path.exists(pdb_file_path):
                    return json.dumps({
                        "success": False,
                        "error": f"PDB文件不存在: {pdb_file_path}"
                    }, ensure_ascii=False)
                
                if not pdb_file_path.endswith('.pdb'):
                    return json.dumps({
                        "success": False,
                        "error": "只支持.pdb格式文件"
                    }, ensure_ascii=False)
                
                # 检查文件大小
                file_size = os.path.getsize(pdb_file_path)
                if file_size > config.LIMITS["foldseek"]["max_pdb_file_size"]:
                    return json.dumps({
                        "success": False,
                        "error": f"文件过大，最大允许{config.LIMITS['foldseek']['max_pdb_file_size'] // (1024*1024)}MB"
                    }, ensure_ascii=False)
            
            # 获取其他参数
            tm_score_threshold = params.get("tm_score_threshold", config.DEFAULT_PARAMS["foldseek"]["tm_score_threshold"])
            max_results = params.get("max_results", config.DEFAULT_PARAMS["foldseek"]["max_results"])
            
            # 参数验证
            if not (config.LIMITS["foldseek"]["min_tm_score_threshold"] <= tm_score_threshold <= config.LIMITS["foldseek"]["max_tm_score_threshold"]):
                return json.dumps({
                    "success": False,
                    "error": f"TM-score阈值必须在{config.LIMITS['foldseek']['min_tm_score_threshold']}到{config.LIMITS['foldseek']['max_tm_score_threshold']}之间"
                }, ensure_ascii=False)
            
            if not (config.LIMITS["foldseek"]["min_max_results"] <= max_results <= config.LIMITS["foldseek"]["max_max_results"]):
                return json.dumps({
                    "success": False,
                    "error": f"最大结果数必须在{config.LIMITS['foldseek']['min_max_results']}到{config.LIMITS['foldseek']['max_max_results']}之间"
                }, ensure_ascii=False)
            
            # 构建API请求数据
            api_data = {
                "tm_score_threshold": tm_score_threshold,
                "max_results": max_results
            }
            
            files = None
            if query_accession:
                # 使用数据库中的蛋白质
                api_data["query_accession"] = query_accession
                response = requests.post(
                    config.API_ENDPOINTS["foldseek"],
                    data=api_data,
                    timeout=config.LIMITS["foldseek"]["timeout"]
                )
            else:
                # 上传PDB文件
                with open(pdb_file_path, 'rb') as f:
                    files = {
                        "pdb_file": (os.path.basename(pdb_file_path), f, "application/octet-stream")
                    }
                    # 注意：requests在files存在时会自动设置Content-Type为multipart/form-data
                    response = requests.post(
                        config.API_ENDPOINTS["foldseek"],
                        data=api_data,
                        files=files,
                        timeout=config.LIMITS["foldseek"]["timeout"]
                    )
            
            response.raise_for_status()
            api_result = response.json()
            
            # 格式化结果
            formatted_result = {
                "success": api_result.get("success", True),
                "message": api_result.get("message", ""),
                "query_identifier": api_result.get("query_identifier", query_accession or os.path.basename(pdb_file_path)),
                "total_hits": api_result.get("total_hits", 0),
                "execution_time": api_result.get("execution_time", 0),
                "cached": api_result.get("cached", False),
                "results": []
            }
            
            # 格式化结果列表
            for item in api_result.get("results", []):
                formatted_result["results"].append({
                    "id": item.get("id"),
                    "accession": item.get("accession"),
                    "species": item.get("species"),
                    "type": item.get("type"),
                    "assembly": item.get("assembly"),
                    "has_pdb": item.get("has_pdb", False),
                    "tm_score": item.get("tm_score"),
                    "probability": item.get("probability"),
                    "rmsd": item.get("rmsd"),
                    "alignment_length": item.get("alignment_length"),
                    "target_length": item.get("target_length")
                })
            
            return json.dumps(formatted_result, ensure_ascii=False)
            
        except requests.exceptions.Timeout:
            error_msg = {
                "success": False,
                "error": f"搜索超时（>{config.LIMITS['foldseek']['timeout']}秒）"
            }
            return json.dumps(error_msg, ensure_ascii=False)
        
        except requests.exceptions.RequestException as e:
            error_msg = {
                "success": False,
                "error": f"API调用失败: {str(e)}"
            }
            return json.dumps(error_msg, ensure_ascii=False)
        
        except Exception as e:
            error_msg = {
                "success": False,
                "error": f"工具执行失败: {str(e)}"
            }
            return json.dumps(error_msg, ensure_ascii=False)

